/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#ifndef EXAMPLES_ARITH_PROGRESSION_H
#define EXAMPLES_ARITH_PROGRESSION_H
#include "kernel_operator.h"

constexpr uint32_t BYTES_ALIGN_SIZE = 32;
namespace MyCustomKernel {
struct VecTiling{
    uint32_t outSize;
};

template <typename T>
class KernelArithProgression {
public:
    __aicore__ inline KernelArithProgression()
    {}
    __aicore__ inline void Init(GM_ADDR dstGm, float firstValue, float diffValue, uint32_t count)
    {
        firstValue_ = firstValue;
        diffValue_ = diffValue;
        count_ = count;
        dst_global.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(dstGm), count_);
        pipe.InitBuffer(outDst, 1, (sizeof(T) * count_ + BYTES_ALIGN_SIZE - 1) / BYTES_ALIGN_SIZE * BYTES_ALIGN_SIZE);
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {}
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<T> dstLocal = outDst.AllocTensor<T>();
        AscendC::ArithProgression<T>(dstLocal, static_cast<T>(firstValue_), static_cast<T>(diffValue_), count_);
        outDst.EnQue<T>(dstLocal);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<T> dstLocal = outDst.DeQue<T>();
        const int32_t BLOCK_NUM = BYTES_ALIGN_SIZE / sizeof(T);
        AscendC::DataCopy(dst_global, dstLocal, (count_ + BLOCK_NUM - 1) / BLOCK_NUM * BLOCK_NUM);
        outDst.FreeTensor(dstLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, 1> outDst;
    AscendC::GlobalTensor<T> dst_global;
    float firstValue_;
    float diffValue_;
    uint32_t count_;
};
}
#endif // EXAMPLES_ARITH_PROGRESSION_H